"""
可视化使用示例：展示如何使用新的可视化工具
"""
import numpy as np
import sys
import os

# 添加路径
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from Squidiff.hic_visualization import visualize_forward_diffusion
from Squidiff import diffusion

# 创建示例数据（模拟预处理后的 Hi-C 矩阵）
print("创建示例 Hi-C 矩阵...")
size = 128
x0 = np.random.exponential(2.0, (size, size)).astype(np.float32)
# 使对称
x0 = (x0 + x0.T) / 2
# 对角线更强
for i in range(size):
    for j in range(size):
        distance = abs(i - j)
        if distance < 20:
            x0[i, j] *= (20 - distance) / 20 * 5

print(f"矩阵形状: {x0.shape}")
print(f"矩阵范围: [{x0.min():.2f}, {x0.max():.2f}]")

# 创建 beta schedule（线性，1000 步）
print("\n创建 beta schedule...")
num_timesteps = 1000
betas = diffusion.get_named_beta_schedule("linear", num_timesteps)
print(f"Beta 范围: [{betas.min():.6f}, {betas.max():.6f}]")

# 可视化前向扩散过程
print("\n可视化前向扩散过程...")
output_dir = "./example_visualizations"
os.makedirs(output_dir, exist_ok=True)

# 示例 1: 基础可视化（16 个时间步）
print("\n示例 1: 基础可视化（均匀采样 16 个时间步）")
visualize_forward_diffusion(
    x0=x0,
    betas=betas,
    num_timesteps=16,
    timestep_strategy="linspace",
    use_mean=False,  # 加随机噪声
    log_display=True,
    signed_display=False,
    per_panel_scale=False,
    save_path=os.path.join(output_dir, "example_1_linspace_16steps.png"),
    suptitle="示例 1: 前向扩散（linspace, 16 步）",
)
print(f"  保存至: {output_dir}/example_1_linspace_16steps.png")

# 示例 2: 只显示均值（无噪声）
print("\n示例 2: 只显示均值轨迹（无随机噪声）")
visualize_forward_diffusion(
    x0=x0,
    betas=betas,
    num_timesteps=16,
    timestep_strategy="linspace",
    use_mean=True,  # 只显示均值
    log_display=True,
    signed_display=False,
    per_panel_scale=False,
    save_path=os.path.join(output_dir, "example_2_mean_only.png"),
    suptitle="示例 2: 前向扩散均值轨迹（无噪声）",
)
print(f"  保存至: {output_dir}/example_2_mean_only.png")

# 示例 3: 更密集的时间步（32 步）
print("\n示例 3: 密集采样（32 个时间步）")
visualize_forward_diffusion(
    x0=x0,
    betas=betas,
    num_timesteps=32,
    timestep_strategy="linspace",
    use_mean=False,
    log_display=True,
    signed_display=False,
    per_panel_scale=False,
    save_path=os.path.join(output_dir, "example_3_dense_32steps.png"),
    suptitle="示例 3: 前向扩散（密集采样，32 步）",
)
print(f"  保存至: {output_dir}/example_3_dense_32steps.png")

# 示例 4: 每子图独立色阶
print("\n示例 4: 每子图独立色阶")
visualize_forward_diffusion(
    x0=x0,
    betas=betas,
    num_timesteps=16,
    timestep_strategy="linspace",
    use_mean=False,
    log_display=True,
    signed_display=False,
    per_panel_scale=True,  # 每子图独立缩放
    save_path=os.path.join(output_dir, "example_4_per_panel_scale.png"),
    suptitle="示例 4: 前向扩散（独立色阶）",
)
print(f"  保存至: {output_dir}/example_4_per_panel_scale.png")

# 示例 5: 只看前半段扩散（t <= 500）
print("\n示例 5: 只看前半段扩散过程")
visualize_forward_diffusion(
    x0=x0,
    betas=betas,
    num_timesteps=16,
    timestep_strategy="linspace",
    max_timestep=500,  # 限制最大 t
    use_mean=False,
    log_display=True,
    signed_display=False,
    per_panel_scale=False,
    save_path=os.path.join(output_dir, "example_5_first_half.png"),
    suptitle="示例 5: 前半段扩散（t ≤ 500）",
)
print(f"  保存至: {output_dir}/example_5_first_half.png")

# 示例 6: 步长采样（每 100 步）
print("\n示例 6: 步长采样（每 100 步）")
visualize_forward_diffusion(
    x0=x0,
    betas=betas,
    num_timesteps=12,
    timestep_strategy="linspace",
    timestep_stride=100,  # 每 100 步选一个
    use_mean=False,
    log_display=True,
    signed_display=False,
    per_panel_scale=False,
    save_path=os.path.join(output_dir, "example_6_stride_100.png"),
    suptitle="示例 6: 前向扩散（步长 100）",
)
print(f"  保存至: {output_dir}/example_6_stride_100.png")

print("\n" + "=" * 80)
print("所有示例已生成！")
print(f"查看结果: {output_dir}/")
print("=" * 80)
print("\n提示：")
print("  - 第一张子图是 t=0（预处理后的数据）")
print("  - 后续子图展示逐步加噪的结果")
print("  - 最后接近 t=999 时应该看起来像随机噪声")
print("  - ᾱ_t 从 ~1.0 衰减到 ~0.0")

